from DOA import DOA
import numpy as np
import multiprocessing
from functools import partial
import pickle
import os
import time

# 实现了多线程仿真，兼容不同DOA方法，不同采样数量和仿真任务
# 通过改变参数生成不同仿真数据进行实验，汇总多线程数据保存至pkl
class DOA_concurrent(DOA):
    def __init__(self):
        super(DOA_concurrent, self).__init__()
        self.default_snr_range = (-10, 30)
        self.default_len_s_range = (1, 1500)
        self.default_len_s_sample_times = 15
        self.default_snr_sample_times = 40
        self.workers = 8
        self.num_tests = 1000
        self.evaluate_data = {'num_tests':1000,
                              }


    def evaluation_single_threading(self, sample_data, index_tuple, doa_method, lock1, eva_task='snr', result_array=None, snr_mask=None, wrong_esti_signal_num=None):
        '''
        1. 不同的仿真任务本质上通过不同的task_data输入，创建不同的仿真数据，DOA方法是固定的；根据不同的任务执行不同的数据生成分支
        2. result_array原先用于存储每一次仿真的RMSE，后改为兼容角度输出结果
        3. 仿真采样类别分为snr, len_s, snr_angle等，DOA方法分为平均采样和阵列DOA，这些都单独写了兼容
        '''
        # sample_index为横轴采样的索引，doa_index是第几个doa的索引
        sample_index, method_index = index_tuple
        sample_times = self.evaluate_data["sample_times"]
        num_tests = self.evaluate_data["num_tests"]


        total_errors = 0
        total_angle = 0
        count_test_signal = 0
        theta_true = None
        for _ in range(num_tests):
            # 根据不同仿真任务进入不同数据生成分支
            if eva_task in ['snr', 'snr_angle']:
                snr = sample_data
                if doa_method.__name__ == "mean_intensity":
                    theta_true = np.random.uniform(-np.pi / 2, np.pi / 2)
                    p, v_x, v_y = self.generate_signal_mean_intensity(theta_true, self.default_len_s, snr)
                    signal = (p, v_x, v_y)
                else:
                    # 如果是snr_angle即固定角度测试，设置两个角度，注意前提是信源数为2
                    fixed_angle = [30, 160] if eva_task == "snr_angle" else None
                    signal, theta_true = self.generate_signal(snr, print_angle=False, len_signal=self.default_len_s, fixed_angle=fixed_angle)
            elif eva_task == "len_s":
                len_s = int(sample_data)
                if doa_method.__name__ == "mean_intensity":
                    theta_true = np.random.uniform(-np.pi / 2, np.pi / 2)
                    p, v_x, v_y = self.generate_signal_mean_intensity(theta_true, len_s, SNR=10)
                    signal = (p, v_x, v_y)
                else:
                    signal, theta_true = self.generate_signal(SNR=10, print_angle=False, len_signal=len_s)

            estimated_theta = doa_method(signal)

            # 平均声强法只输出一个方位角信息，和多信号源DOA分开处理
            if not doa_method.__name__ == "mean_intensity":
                estimated_theta = np.sort(estimated_theta)
                # 搜索角度和真值不对应暂时先不考虑，不参与统计
                if len(estimated_theta) != len(theta_true):
                    lock1.acquire()
                    wrong_esti_signal_num.value += 1
                    lock1.release()
                    continue
                # 遍历本次测试所有估计信号源角度
                for i, theta in enumerate(estimated_theta):
                    count_test_signal += 1
                    # 输出角度结果和输出误差的统计方法不同
                    if eva_task in ['snr', 'len_s']:total_errors += (theta - theta_true[i]) ** 2
                    elif eva_task in ['snr_angle']:total_angle += theta
            else:  # 平均声强法的处理，一次测试只有一个方位角
                if estimated_theta is not None:
                    count_test_signal += 1
                    total_errors += (theta_true - estimated_theta) ** 2
                else:
                    continue

        if count_test_signal > 0:
            total_result = 0  # 用于计算均值
            if eva_task in ['snr', 'len_s']: total_result = total_errors
            elif eva_task in ['snr_angle']: total_result = total_angle
            result_array[sample_index + method_index * sample_times] = total_result / count_test_signal
        else:
            snr_mask[sample_index + method_index*sample_times] = False
            if eva_task == 'snr':
                print(f"sample_test_num:{sample_index + method_index*sample_times},count zero!  {eva_task}={sample_data:.4f}")
            else:
                print("undefined, count zero")

    # 对N个横轴采样，M个方法进行并发1000次仿真，线程池中为N*M个任务
    def evaluation_multiThread(self, doa_methods, eva_task="snr", sample_range_input=None, sample_times_input=None):
        '''
        1. doa_methods中输入多个DOA方法，以循环的方式和横轴采样一起放入all_tasks中，放入二者的index以从数组中索引写入
        2. ...
        '''
        from tools import par_run, par_wrapper
        # 根据仿真任务初始化对应参数
        if eva_task in ["snr", 'snr_angle']:
            sample_range = self.default_snr_range  # (-10, 30)
            sample_times = self.default_snr_sample_times  # 40
        elif eva_task == "len_s":
            sample_range = self.default_len_s_range  # (100, 1500)
            sample_times = self.default_len_s_sample_times  # 20
        else:
            sample_range = None
            sample_times = None
            print("no task for default sample range")
        # sample_range根据仿真任务有默认值，可以通过参数强制修改
        if sample_range_input is not None: sample_range = sample_range_input
        if sample_times_input is not None: sample_times = sample_times_input

        all_tasks = []
        # lock = multiprocessing.Lock()
        methods_num = len(doa_methods)
        rmse_array = np.zeros(sample_times*methods_num)
        rmse_mask = np.ones(sample_times*methods_num, dtype=bool)
        wrong_esti_signal_num = 0

        sample_array = np.linspace(sample_range[0], sample_range[1], sample_times)  # 等间隔采样
        self.evaluate_data['num_methods']=methods_num
        self.evaluate_data['sample_times'] = sample_times
        self.evaluate_data['sample_array'] = sample_array
        # ==================
        # 配置进程间数据
        manager = multiprocessing.Manager()
        shared_rmse_mask = manager.Array('b', rmse_mask)
        shared_rmse_array = manager.Array('f', rmse_array)
        shared_wrong_esti_signal_num = manager.Value("i", 0)
        lock = manager.Lock()
        # 填充Task_data
        for sample_index, sample_data in enumerate(sample_array):  # 遍历所有横轴采样
            for doa_index, doa_method in enumerate(doa_methods):  # 遍历所有方法
                # sample_index, doa_index: 第几个横轴采样和DOA方法，形成总的任务数量，传入各自索引构建data数组
                all_tasks.append((sample_data, (sample_index, doa_index), doa_method, lock, eva_task, shared_rmse_array, shared_rmse_mask, shared_wrong_esti_signal_num))
        # 执行并行处理
        par_run(partial(par_wrapper, func=self.evaluation_single_threading), all_tasks, b_thread=False, n_par=self.workers, desc="[testing SNR]")

        # 解析数据
        rmse_mask = np.array([bool(item) for item in shared_rmse_mask])
        rmse_array = np.array([item for item in shared_rmse_array])
        print(f"error signals number: {wrong_esti_signal_num}！")
        manager.shutdown()

        # 保存数据为Pickle
        output_dir = eva_task
        os.makedirs(output_dir, exist_ok=True)
        for method_i, method in enumerate(doa_methods):
            method_rmse_mask = rmse_mask[method_i * sample_times:(method_i + 1) * sample_times]
            method_rmse = rmse_array[method_i * sample_times:(method_i + 1) * sample_times]
            method_sample_array = sample_array[method_rmse_mask]  # sample_array在串行外部创建，通过method对应mask可以得到当前method sample array，每个方法的初始sample array相同
            data = np.vstack((method_sample_array, method_rmse))
            pkl_name = f"{method.__name__}_{self.num_tests}test_{sample_times}sample.pkl"
            with open(output_dir+"/"+pkl_name, "wb") as f:
                pickle.dump(data, f)

        # return SNR_array, rmse_array



if __name__ == "__main__":
    doa_concurrent = DOA_concurrent()
    # "len_s" "snr" "snr_angle"
    # methods = (doa_concurrent.music, doa_concurrent.root_music, doa_concurrent.esprit)
    methods = (doa_concurrent.mean_intensity,)
    start_time = time.time()
    # snr和len_s方法在不同角度下测试RMSE，snr_angle任务在相同下测试估计效果
    doa_concurrent.evaluation_multiThread(methods, eva_task='len_s')
    end_time = time.time()
    print(f"程序执行时间为：{end_time - start_time} 秒")

